#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <pthread.h>
#include "tbf.h"

typedef struct {
	int token;
	int cps;
	int burst;
	pthread_mutex_t mut;
	pthread_cond_t cond;
}tbf_t;

static pthread_mutex_t lib_mut = PTHREAD_MUTEX_INITIALIZER;
static tbf_t *libs[MAXTBF];
static pthread_once_t once_control = PTHREAD_ONCE_INIT;
static pthread_t tid;

static void moduler_load();

static int __get_free_pos(void)
{
	int i;

	pthread_mutex_lock(&lib_mut);
	for (i = 0; i < MAXTBF; i++) {
		if (libs[i] == NULL) {
			pthread_mutex_unlock(&lib_mut);
			return i;
		}
	}

	pthread_mutex_unlock(&lib_mut);
	
	return -1;
}

int tbf_init(int cps, int burst)
{
	tbf_t *t = NULL;
	int pos;

	// 当第一次调用初始化令牌桶，启动异步线程
	pthread_once(&once_control, moduler_load);

	// 创建令牌桶
	// 开辟存储空间
	t = malloc(sizeof(tbf_t));
	if (NULL == t)
		return -1;
	// 成员初始化
	t->token = 0;
	t->cps = cps;
	t->burst = burst;
	pthread_mutex_init(&t->mut, NULL);
	pthread_cond_init(&t->cond, NULL);
	// 存放到令牌桶库中
	pos = __get_free_pos();
	if (-1 == pos) {
		free(t);
		t = NULL;
		return -1;
	}
	libs[pos] = t;

	return pos;
}

static void *thr_job(void *s)
{
	int i;

	// 周期性遍历令牌桶库中的每一个令牌桶，令牌积攒
	while (1) {
		pthread_mutex_lock(&lib_mut);
		for (i = 0; i < MAXTBF; i++) {
			if (libs[i]) {
				// 积攒令牌
				pthread_mutex_lock(&libs[i]->mut);
				libs[i]->token += libs[i]->cps;
				if (libs[i]->token > libs[i]->burst) {
					libs[i]->token = libs[i]->burst;
				}
				pthread_cond_broadcast(&libs[i]->cond);
				pthread_mutex_unlock(&libs[i]->mut);
			}
		}
		pthread_mutex_unlock(&lib_mut);
		sleep(1);
	}
}

static void moduler_load()
{
	int err;
	// 启动线程
	err = pthread_create(&tid, NULL, thr_job, NULL);
	if (err) {
		fprintf(stderr, "pthread_create():%s\n", strerror(err));
	}
}

int tbf_fetch_token(int td, int ntoken)
{
	int get_cnt = ntoken;

	if (!(td >= 0 && td < MAXTBF))
		return -1;
	pthread_mutex_lock(&libs[td]->mut);
	while (libs[td]->token <= 0) {
		pthread_cond_wait(&libs[td]->cond, &libs[td]->mut);
	}
	if (ntoken > libs[td]->token) {
		get_cnt = libs[td]->token;	
	}
	libs[td]->token -= get_cnt;
	pthread_mutex_unlock(&libs[td]->mut);

	return get_cnt;
}

int tbf_return_token(int td, int ntoken)
{
	if (!(td >= 0 && td < MAXTBF))
		return -1;
	pthread_mutex_lock(&libs[td]->mut);
	libs[td]->token += ntoken;	
	if (libs[td]->token > libs[td]->burst)
		libs[td]->token = libs[td]->burst;
	pthread_cond_broadcast(&libs[td]->cond);
	pthread_mutex_unlock(&libs[td]->mut);

	return 0;
}

void tbf_destroy(int td)
{
	pthread_mutex_lock(&lib_mut);
	pthread_mutex_destroy(&libs[td]->mut);
	pthread_cond_destroy(&libs[td]->cond);
	free(libs[td]);
	libs[td] = NULL;
	pthread_mutex_unlock(&lib_mut);
}

void tbf_destroy_all(void)
{
	int i;

	pthread_cancel(tid);
	for (i = 0; i < MAXTBF; i++) {
		if (libs[i]) {
			tbf_destroy(i);
		}
	}
	pthread_mutex_destroy(&lib_mut);
}


